#!  /usr/bin/env python
# coding=utf-8

import rospy
import math
from pyquaternion import Quaternion
import numpy as np
import json
import matplotlib.pyplot as plt
from ti5arm.msg import ti5_polyJC_Lists
pi = math.pi

from ti5_fkik import ti5FK


def jDeg2q(seq,json_name="armJ_handeyeXYZwxyz",ori_dir="/home/wheeltec/Desktop/ros_cam/aruco_hand_eye_cli_file.txt", 
           Q_save_dir = '/home/wheeltec/Desktop/ros_cam/hand_eyeQ.txt',
           xyz_save_dir = '/home/wheeltec/Desktop/hand_eyeXYZ.txt',
           xyzQ_save_dir = '/home/wheeltec/Desktop/ros_cam/hand_eye_XYZ_Q.txt'):
    """关节角度(单位：度)转为位姿XYZwxyz(XYZ单位：m)并写入文件
    
    input：

        seq: 写入新文件的数据顺序 ('wxyz', 'xyzw', 'XYZwxyz', 'XYZxyzw')

        json_name: dir文件第一行的name

        ori_dir: 要打开的文件
        Q_save_dir:四元数保存位置 seq='wxyz'或 'xyzw' 时有效
        xyz_save_dir:三维坐标保存位置  seq='wxyz'或 'xyzw' 时有效
        xyzQ_save_dir:三维坐标及四元数保位置  seq='XYZwxyz'或 'XYZxyzw' 时有效
    """
    try:
        with open(ori_dir, mode='r') as Jfile:
            Jdata = Jfile.readlines()
            pass
    except Exception as e:
        print(e)
        print('open file hand_eyeJ failed!')
        return

    groups_jdeg = (len(Jdata)-1)
    Jdeg = np.zeros((groups_jdeg, 6), dtype=np.float) # type:np.ndarray
    h_eQ = np.zeros((groups_jdeg, 4), dtype=np.float) # type:np.ndarray
    eefXYZ = np.zeros((groups_jdeg, 3), dtype=np.float) # type:np.ndarray
    try:
        head = json.loads(Jdata[0])
        readname = head['name']
      
    except Exception as e:
        print('json loads msg_head failed!')
        print(e)
        readname = ''
        return

    if(readname==json_name):# 
        print("open file!")
        print("groups of Jdeg:",groups_jdeg)
        for i in range(0,groups_jdeg):
            Jdata_split = Jdata[1+i].split(',')
            # print(Jdata_split)
            Jdeg[i,0] = float(Jdata_split[0])
            Jdeg[i,1] = float(Jdata_split[1])
            Jdeg[i,2] = float(Jdata_split[2])
            Jdeg[i,3] = float(Jdata_split[3])
            Jdeg[i,4] = float(Jdata_split[4])
            Jdeg[i,5] = float(Jdata_split[5])

            tF = ti5FK(Jdeg[i]*pi/180)
            eefXYZ[i,0] = tF[0,3]
            eefXYZ[i,1] = tF[1,3]
            eefXYZ[i,2] = tF[2,3]
            q = Quaternion(matrix = tF[0:3,0:3])
            if(seq=='wxyz' or seq=='XYZwxyz'):
                h_eQ[i,0] = q.w 
                h_eQ[i,1] = q.x 
                h_eQ[i,2] = q.y 
                h_eQ[i,3] = q.z 
            elif(seq=='xyzw' or seq=='XYZxyzw'):
                h_eQ[i,0] = q.x 
                h_eQ[i,1] = q.y 
                h_eQ[i,2] = q.z 
                h_eQ[i,3] = q.w 
            else:
                print('seq error!')
                return
            pass
        if(seq=='wxyz' or seq=='xyzw'):
            with open(Q_save_dir, mode='w') as f:
                head = {'name':'hand_eyeQ', 'writeTime':rospy.Time.now().to_sec(), 
                        'seq':seq,'shape':[groups_jdeg,4]}
                msg_head = json.dumps(head)
                f.write(msg_head)
            
                for i in range(0, groups_jdeg):
                    w_msg = "\r\n{}, {}, {}, {}".format(h_eQ[i,0],
                                                        h_eQ[i,1],
                                                        h_eQ[i,2],
                                                        h_eQ[i,3])
                    f.write(w_msg)
            print('hand_eyeQ write to %s \n'%(Q_save_dir))
            
            with open(xyz_save_dir, mode='w') as f:
                head = {'name':'hand_eyeXYZ', 'writeTime':rospy.Time.now().to_sec(), 
                        'seq':'XYZ','shape':[groups_jdeg,3]}
                msg_head = json.dumps(head)
                f.write(msg_head)
            
                for i in range(0, groups_jdeg):
                    w_msg = "\r\n{}, {}, {}".format(eefXYZ[i,0],
                                                    eefXYZ[i,1],
                                                    eefXYZ[i,2])
                    f.write(w_msg)
            print('hand_eyeXYZ write to %s'%(xyz_save_dir))

            pass
        elif(seq=='XYZwxyz' or seq=='XYZxyzw'):
            with open(xyzQ_save_dir, mode='w') as f:
                head = {'name':'hand_eye_XYZ_Q', 'writeTime':rospy.Time.now().to_sec(), 
                        'seq':seq,'shape':[groups_jdeg,7]}
                msg_head = json.dumps(head)
                f.write(msg_head)
            
                for i in range(0, groups_jdeg):
                    w_msg = "\r\n{}, {}, {}, {}, {}, {}, {}".format(eefXYZ[i,0]/1000,
                                                                    eefXYZ[i,1]/1000,
                                                                    eefXYZ[i,2]/1000,
                                                                    h_eQ[i,0],
                                                                    h_eQ[i,1],
                                                                    h_eQ[i,2],
                                                                    h_eQ[i,3])
                    f.write(w_msg)
            print('hand_eye_XYZ_Q write to %s'%(xyzQ_save_dir))
            pass
        else:
            print('seq error!')
            return
            pass
    else:
        print('name is not {}!'.format(json_name),readname)
       
    pass

def split_arucoPose(json_name, ori_dir, 
                    save_dir = '/home/wheeltec/Desktop/ros_cam/arucoPose.txt'):
    """从给定armJ_handeyeXYZwxyz文件分离出arucoPose(XYZwxyz)并写入新文件
    
    input：
    
        json_name: dir文件第一行的name

        ori_dir: 要打开的文件位置
        save_dir:要保存的位置
    """
    try:
        with open(ori_dir, mode='r') as Jfile:
            Jdata = Jfile.readlines()
            pass
    except Exception as e:
        print(e)
        print('open file hand_eyeJ failed!')
        return
    
    groups_jdeg = (len(Jdata)-1)
    XYZwxyz = np.zeros((groups_jdeg, 7), dtype=np.float) # type:np.ndarray
    try:
        head = json.loads(Jdata[0])
        readname = head['name'] 
    except Exception as e:
        print('json loads msg_head failed!')
        print(e)
        readname = ''
        return

    if(readname==json_name):
        print("open file!")
        print("groups of Jdeg:",groups_jdeg)
        for i in range(0,groups_jdeg):
            Jdata_split = Jdata[1+i].split(',')
            # print(Jdata_split)
            XYZwxyz[i,0] = float(Jdata_split[6])
            XYZwxyz[i,1] = float(Jdata_split[7])
            XYZwxyz[i,2] = float(Jdata_split[8])
            XYZwxyz[i,3] = float(Jdata_split[9])
            XYZwxyz[i,4] = float(Jdata_split[10])
            XYZwxyz[i,5] = float(Jdata_split[11])
            XYZwxyz[i,6] = float(Jdata_split[12])
            pass
        
        with open(save_dir, mode='w') as f:
            head = {'name':'hand_eyeQ', 'writeTime':rospy.Time.now().to_sec(), 
                    'seq':'XYZwxyz','shape':[groups_jdeg,4]}
            msg_head = json.dumps(head)
            f.write(msg_head)
        
            for i in range(0, groups_jdeg):
                w_msg = "\r\n{}, {}, {}, {}, {}, {}, {}".format(XYZwxyz[i,0],
                                                                XYZwxyz[i,1],
                                                                XYZwxyz[i,2],
                                                                XYZwxyz[i,3],
                                                                XYZwxyz[i,4],
                                                                XYZwxyz[i,5],
                                                                XYZwxyz[i,6])
                f.write(w_msg)
        print('arucoPose write to %s'%(save_dir))


    pass


def q2Rot(q,seq='wxyz'):
    """
        四元素转旋转矩阵
    input:
        q:w,x,y,z
        seq:四元素沙顺序
    """

    if seq=='wxyz':
        w = q[0]
        x = q[1]
        y = q[2]
        z = q[3]
    elif seq=='xyzw':
        w = q[3]
        x = q[0]
        y = q[1]
        z = q[2]
    else:
        print('input seq err,no seq:{}'.format(seq))
        return 
    Rq = np.array([[1-2*y*y-2*z*z, 2*x*y-2*w*z,   2*x*z+2*w*y],
                  [2*x*y+2*w*z,   1-2*x*x-2*z*z, 2*y*z-2*w*x],
                  [2*x*z-2*w*y,   2*y*z+2*w*x,   1-2*x*x-2*y*y]])


    return Rq
    pass

def PvRot2T(p,rot):
    """
        平移向量与旋转矩阵转变换矩阵
    input:
        p:平移向量(XYZ)
        rot:旋转矩阵
    """
    matT = np.zeros((4,4))

    matT[0:3,0:3] = rot
    matT[0:3,3] = np.array([p[0],p[1],p[2]])
    matT[3,:] = np.array([0,0,0,1])

    return matT
    pass


def verify_handeye(b2c_XYZwxyz,Mat_c2B):
    """
        验证手眼标定结果
    input:
        b2c_XYZwxyz:标定板到相机的变换参数np.ndarray
        Mat_c2B:相机到基座的变换矩阵np.ndarray
    """

    # 标定板到相机的变换矩阵
    Mat_b2c = PvRot2T(p=b2c_XYZwxyz[0:3],rot=q2Rot(q=b2c_XYZwxyz[3:7]))

    # 标定板到基座的变换矩阵
    Mat_b2B = np.dot(Mat_c2B,Mat_b2c)
    print("Mat_b2B:")
    print(np.around(Mat_b2B,6))
    pass



from ti5GrabFruits import sub_motorStat
def ti5_getCJdeg():
    """堵塞20ms至收到关节消息/deg,返回的关节角为二维数组
    """
    for i in range(0,10):
        cJdeg = sub_motorStat()
        if  type(cJdeg) != type(-1):
            return cJdeg
    return -1
    pass


from ti5arm.msg import ti5_targetPos

def hand_eye_cli_manual():

    print("开始标定...")


    hand_eye_cmd_pub = rospy.Publisher("hand_eye_cmd",ti5_targetPos,queue_size=1)
    hand_eye_cmd = ti5_targetPos()
    hand_eye_cmd.cmdList = "0" 

    hand_eye_cmd.joint1RadList.append(0)
    hand_eye_cmd.joint2RadList.append(0)
    hand_eye_cmd.joint3RadList.append(0)
    hand_eye_cmd.joint4RadList.append(0)
    hand_eye_cmd.joint5RadList.append(0)
    hand_eye_cmd.joint6RadList.append(0)

    while True:
        print("输入'q'保存数据并结束标定，输入'w'记录一组数据 (输入字符串需要包括引号)")
        cmd_in = input("请输入选项：")
        if cmd_in == "q":
            hand_eye_cmd.cmdList = "1"
            hand_eye_cmd.header.stamp = rospy.Time().now()
            hand_eye_cmd_pub.publish(hand_eye_cmd)
            print("正在退出...")
            return
        elif cmd_in == "w":
            hand_eye_cmd.cmdList = "2"
            print("cmd 2")
            arm_cjDeg = ti5_getCJdeg()
            if type(arm_cjDeg) != type(-1):


                hand_eye_cmd.joint1RadList[0] = arm_cjDeg[0,0]
                hand_eye_cmd.joint2RadList[0] = arm_cjDeg[0,1]
                hand_eye_cmd.joint3RadList[0] = arm_cjDeg[0,2]
                hand_eye_cmd.joint4RadList[0] = arm_cjDeg[0,3]
                hand_eye_cmd.joint5RadList[0] = arm_cjDeg[0,4]
                hand_eye_cmd.joint6RadList[0] = arm_cjDeg[0,5]

                print('arm_cjDeg',arm_cjDeg)


                hand_eye_cmd.header.stamp = rospy.Time().now()
                hand_eye_cmd_pub.publish(hand_eye_cmd)
            else:
                print("cope cmd 'w' failed, please retry")
             
        else :
 
            print(" unknown cmd ! please retry ")
        pass

    pass




if  __name__== "__main__":
    rospy.init_node("hand_eye_cli_node")

    #******标定******
    # hand_eye_cli_manual()
    #***************

    #*****求解*******
    # jDeg2q( seq="XYZwxyz",
    #         json_name="armJ_handeyeXYZwxyz",
    #         ori_dir="/home/wheeltec/Desktop/ros_cam/handeyearmJ_2443.txt", 
    #         Q_save_dir = '/home/wheeltec/Desktop/ros_cam/hand_eyeQ_2443.txt',
    #         xyz_save_dir = '/home/wheeltec/Desktop/hand_eyeXYZ_2443.txt',
    #         xyzQ_save_dir = '/home/wheeltec/Desktop/ros_cam/hand_eye_XYZ_Q_2443.txt')
    # split_arucoPose(json_name="armJ_handeyeXYZwxyz",
    #                 ori_dir="/home/wheeltec/Desktop/ros_cam/aruco_hand_eye_cli_file124331943.txt",
    #                 save_dir = '/home/wheeltec/Desktop/ros_cam/arucoPose2443.txt')
    #***************

    #*****验证*******
    # verify_handeye(b2c_XYZwxyz=np.array([0.156289,0.010014,0.72105169,
    #                       -0.010411,0.9998756,-0.0017207,-0.0117147]),
    #                       Mat_c2B=np.array([[0.0088948,-0.998813,-0.047888,0.6084318],
    #                                         [-0.99959,-0.00757824,-0.027604,0.6556383],
    #                                         [0.0272084,0.04811393,-0.9984712,0.7083566],
    #                                         [0,0,0,1]]))
    # verify_handeye(b2c_XYZwxyz=np.array([0.156289,0.010014,0.72105169,-0.010411,0.9998756,-0.0017207,-0.0117147]),
    #             Mat_c2B=np.array([[0.9998654259299602, 0.01477343174020516, 0.007132723499234949, -0.06918466743426109],
    #                             [0.01413442789919221, -0.9964847433944513, 0.08257344688250037, -0.5220956082144669],
    #                             [0.008327543326910786, -0.08246151767165137, -0.9965594563925557, 0.6975328049459156],
    #                             [0, 0, 0, 1]]))
    #***************
    pass



